package org.farm2.service.utils;

import com.farm2.llm.client.inter.Farm2EmbedClientInter;
import com.farm2.llm.client.inter.Farm2LlmClientInter;
import com.farm2.lucene.Farm2LuceneFace;
import com.farm2.lucene.query.DocumentQuery;
import com.farm2.lucene.query.rule.DocumentRuleByVector;
import com.farm2.lucene.result.DocumentResult;
import com.farm2.lucene.utils.FarmLuceneUtils;
import org.farm2.service.lucene.doc.enums.DocumentTitleEnum;
import org.farm2.tools.db.ResultDataHandle;
import org.farm2.tools.web.dto.QuotingKnowDto;
import org.farm2.tools.web.dto.SendMsgDto;
import org.farm2.tools.caches.FarmCacheKeys;
import org.farm2.tools.caches.FarmCaches;
import org.farm2.tools.db.commons.FarmUUID;

import java.util.Map;

/**
 * 缓存对话消息
 */
public class FarmAIMsgCacheUtils {

    public static String cacheMsg(SendMsgDto msg) {
        String talkKey = FarmUUID.getUUID32();
        msg.setTalkKey(talkKey);
        FarmCaches.getInstance().putCacheData(talkKey, msg, FarmCacheKeys.TALK_MSG);
        return talkKey;
    }

    public static SendMsgDto getMsg(String talkKey, Farm2LlmClientInter talkClient, Farm2EmbedClientInter embedClient) {
        SendMsgDto msg = FarmCaches.getInstance().getCacheData(talkKey, FarmCacheKeys.TALK_MSG);
        if ("all".equals(msg.getSearchKnow()) && embedClient != null) {
            //检索知识库
            DocumentQuery query = new DocumentQuery();
            query.addRule(new DocumentRuleByVector(DocumentTitleEnum.embedding_vector.name(), embedClient.getEmbedding(msg.getMsg()), 8));
            DocumentResult result = Farm2LuceneFace.getSearchService(FarmLuceneUtils.getLuceneEmbeddingFilePath()).search(query);
            result.runDataHandle(new ResultDataHandle() {
                @Override
                public void handle(Map<String, Object> row) {
                    String text = (String) row.get(DocumentTitleEnum.text.name());
                    String knowid = (String) row.get(DocumentTitleEnum.knowId.name());
                    String title = (String) row.get(DocumentTitleEnum.title.name());
                    String fileid = (String) row.get(DocumentTitleEnum.fileId.name());
                    msg.getQuotingKnows().add(new QuotingKnowDto(knowid, fileid, title, text));
                }
            });
        }
        for (String key : msg.getHisTalkKeys()) {
            SendMsgDto hismsg = FarmCaches.getInstance().getCacheData(key, FarmCacheKeys.TALK_MSG);
            msg.getHisTalks().add(hismsg);
        }
        return msg;
    }
}
